/**
 * 项目入口文件
 * @description 这是项目的主入口文件
 * @author RAG Demo Project
 * @version 1.0.0
 * 我们创建一个 SimpleRag 类，主要有四个方法。
 * 1.其中 initialize 是异步初始化方法，vectra 向量数据库通过 new LocalIndex 来创建本地数据索引，在构造函数里，我们可以指定一个 localIndex 的存放地址，我们默认让它存放于 .vectra 目录下，实际上它是以一个简单的 JSON 文本格式存放的。
 * 2.接着是 add 方法，它是用来创建文本内容索引并保存到向量数据库中的，我们可以把任何文本数据通过调用 add 来创建内容索引。
 * 3.我们通过调用 db.deleteItem 方法来批量删除包含指定 id 的记录。
 * 4.最后是 query 方法，它是用来查询向量数据库中与指定文本最相似的内容的。
 * 这样我们就完成了简单的 RAG 工具，它能基于文本内容创建数据，给大模型使用。
 */
import path from "node:path";
import { LocalIndex } from "vectra";
import { getVector, getEmbeddings } from "./src/utils";

export class SimpleRag {
  private db: LocalIndex | undefined;
  private indexPath: string;

  constructor(indexPath: string = ".vectra") {
    this.indexPath = path.join(__dirname, "..", indexPath);
  }

  get avaliable() {
    return this.db !== undefined;
  }

  async initialize() {
    const index = new LocalIndex(this.indexPath);
    if (!(await index.isIndexCreated())) {
      await index.createIndex();
    }
    this.db = index;
  }

  async add(text: string) {
    if (!this.avaliable) throw new Error("RAG is not initialized");

    const embeddings = await getEmbeddings(text);
    const res: (Awaited<ReturnType<LocalIndex["insertItem"]>> | undefined)[] =
      [];
    for (const embedding of embeddings) {
      res.push(await this.db?.insertItem(embedding));
    }
    return res.filter((item) => item).map((item) => ({ id: item!.id }));
  }

  async del(items: { id: string } | { id: string }[]) {
    if (!Array.isArray(items)) items = [items];
    if (!this.avaliable) throw new Error("RAG is not initialized");
    // 这地方同样也不能并行删除
    // return await Promise.all(items.map(async (item) => this.db?.deleteItem(item.id)));
    const res: { id: string }[] = [];
    for (const item of items) {
      await this.db?.deleteItem(item.id);
      res.push({ id: item.id });
    }
    return res;
  }

  async query(query: string, topK: number = 5) {
    if (!this.avaliable) throw new Error("RAG is not initialized");
    const vector = await getVector(query);
    const result = await this.db?.queryItems(vector, query, topK);
    return result?.map(({ item, score }) => ({
      text: item.metadata.text,
      query,
      simularity: score,
      id: item.id,
    }));
  }
}
